"""
ETF分析智能体模块 - 提供ETF分析功能
"""

from typing import Dict, TypedDict, List, Union, Any
from langgraph.graph import StateGraph, END
import json
import datetime
import logging
import pandas as pd

from FDAS.agents.base import BaseAgent, BaseAgentState
from FDAS.data import get_asset
from FDAS.data.assets.base import ASSET_TYPE_ETF

# 自定义JSON编码器，处理日期类型
class DateTimeEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, (datetime.date, datetime.datetime)):
            return obj.isoformat()
        return super(DateTimeEncoder, self).default(obj)

# 定义ETF智能体状态
class ETFAgentState(BaseAgentState):
    symbol: str = ""
    analysis_type: str = ""
    etf_data: Dict = {}

class ETFAnalysisAgent(BaseAgent):
    """ETF分析智能体"""
    
    def _build_workflow(self) -> StateGraph:
        """构建智能体工作流"""
        workflow = StateGraph(ETFAgentState)
        
        # 添加节点
        workflow.add_node("parse_query", self._parse_query)
        workflow.add_node("retrieve_data", self._retrieve_data)
        workflow.add_node("analyze_data", self._analyze_data)
        
        # 设置节点连接关系
        workflow.set_entry_point("parse_query")
        workflow.add_edge("parse_query", "retrieve_data")
        workflow.add_edge("retrieve_data", "analyze_data")
        workflow.add_edge("analyze_data", END)
        
        return workflow
    
    def _parse_query(self, state: ETFAgentState) -> Dict:
        """解析用户查询并提取ETF代码和分析类型"""
        try:
            response = self.llm.invoke(
                "从以下ETF分析请求中提取信息：\n" + state['user_query'] +
                "\n只输出JSON格式，不要任何解释，格式为：\n" +
                '{"symbol": "ETF代码", "analysis_type": "表现分析|跟踪分析|折溢价分析|持仓分析|综合分析"}\n' +
                '确保：1.使用双引号 2.无注释 3.代码通常为6位数字'
            )
            
            # 解析JSON
            result = self._clean_json_response(response.content)
            
            # 数据校验
            result['symbol'] = str(result.get('symbol', '510300')).strip().zfill(6)
            analysis_type = str(result.get('analysis_type', '表现分析')).lower()
            result['analysis_type'] = {
                'performance': '表现分析',
                'tracking': '跟踪分析',
                'premium_discount': '折溢价分析',
                'holdings': '持仓分析',
                'comprehensive': '综合分析',
                '表现': '表现分析',
                '跟踪': '跟踪分析',
                '折溢价': '折溢价分析',
                '持仓': '持仓分析',
                '综合': '综合分析'
            }.get(analysis_type, '表现分析')
            
            return {
                "symbol": result['symbol'],
                "analysis_type": result['analysis_type'],
                "parsed_query": json.dumps(result, ensure_ascii=False, cls=DateTimeEncoder),
                "user_query": state['user_query']
            }
                
        except Exception as e:
            self.logger.error(f"查询解析失败: {str(e)}")
            # 默认返回
            return {
                "symbol": "510300",
                "analysis_type": "表现分析",
                "parsed_query": f"解析失败：{str(e)}",
                "user_query": state['user_query']
            }
    
    def _retrieve_data(self, state: ETFAgentState) -> Dict:
        """获取相关数据"""
        try:
            symbol = state['symbol']
            analysis_type = state['analysis_type']
            
            # 获取ETF对象
            etf = get_asset(symbol, ASSET_TYPE_ETF)
            
            # 基础数据 - 所有分析类型都需要
            basic_info = etf.basic_info
            
            # 根据分析类型获取不同数据
            if analysis_type == '表现分析':
                data = {
                    "basic_info": basic_info,
                    "nav": etf.get_nav().to_dict('records'),
                    "performance": etf.get_performance(),
                    "history": etf.get_history(period="daily", limit=60).to_dict('records')
                }
            elif analysis_type == '跟踪分析':
                data = {
                    "basic_info": basic_info,
                    "nav": etf.get_nav().to_dict('records'),
                    "tracking_error": etf.get_tracking_error(),
                    "history": etf.get_history(period="daily", limit=90).to_dict('records')
                }
            elif analysis_type == '折溢价分析':
                data = {
                    "basic_info": basic_info,
                    "premium_discount": etf.get_premium_discount(),
                    "history": etf.get_history(period="daily", limit=30).to_dict('records')
                }
            elif analysis_type == '持仓分析':
                data = {
                    "basic_info": basic_info,
                    "portfolio": etf.get_portfolio(),
                    "nav": etf.get_nav().to_dict('records')
                }
            elif analysis_type == '综合分析':
                data = {
                    "basic_info": basic_info,
                    "nav": etf.get_nav().to_dict('records'),
                    "portfolio": etf.get_portfolio(),
                    "performance": etf.get_performance(),
                    "tracking_error": etf.get_tracking_error(),
                    "premium_discount": etf.get_premium_discount(),
                    "history": etf.get_history(period="daily", limit=60).to_dict('records')
                }
            else:
                data = {"error": f"暂不支持的分析类型: {analysis_type}"}
            
            return {
                **state,  # 保留现有状态
                "etf_data": data
            }
                
        except Exception as e:
            self.logger.error(f"数据获取失败: {str(e)}")
            return {
                **state,
                "etf_data": {"error": f"数据获取失败: {str(e)}"}
            }
    
    def _analyze_data(self, state: ETFAgentState) -> Dict:
        """执行分析并生成报告"""
        try:
            # 检查数据
            etf_data = state.get('etf_data', {})
            
            if not etf_data:
                return {**state, "analysis_result": self._format_error_response("未获取到有效数据")}
                
            if 'error' in etf_data:
                return {**state, "analysis_result": self._format_error_response(etf_data['error'])}
            
            # 构建分析提示
            analysis_type = state['analysis_type']
            
            # 根据分析类型构建不同的提示
            if analysis_type == '表现分析':
                prompt = f"""
                请基于以下数据对ETF进行表现分析：
                
                基本信息：{json.dumps(etf_data.get('basic_info', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                净值数据：{json.dumps(etf_data.get('nav', {})[:10], ensure_ascii=False, cls=DateTimeEncoder)}
                
                业绩指标：{json.dumps(etf_data.get('performance', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                历史数据：{json.dumps(etf_data.get('history', {})[:5], ensure_ascii=False, cls=DateTimeEncoder)}
                
                用户需求：{state['user_query']}
                
                请生成一份专业的ETF表现分析报告，包含以下要素：
                1. 短期、中期、长期收益表现
                2. 与基准指数对比
                3. 与同类ETF对比
                4. 交易流动性分析
                5. 波动性分析
                6. 投资建议
                
                使用Markdown格式输出，确保专业、客观、有深度。
                """
            elif analysis_type == '跟踪分析':
                prompt = f"""
                请基于以下数据对ETF进行跟踪分析：
                
                基本信息：{json.dumps(etf_data.get('basic_info', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                净值数据：{json.dumps(etf_data.get('nav', {})[:10], ensure_ascii=False, cls=DateTimeEncoder)}
                
                跟踪误差：{json.dumps(etf_data.get('tracking_error', 0.0), ensure_ascii=False, cls=DateTimeEncoder)}
                
                历史数据：{json.dumps(etf_data.get('history', {})[:5], ensure_ascii=False, cls=DateTimeEncoder)}
                
                用户需求：{state['user_query']}
                
                请生成一份专业的ETF跟踪分析报告，包含以下要素：
                1. 跟踪误差分析
                2. 跟踪效率评估
                3. 跟踪偏离原因分析
                4. 复制策略推测
                5. 与同类ETF跟踪效果对比
                6. 投资建议
                
                使用Markdown格式输出，确保专业、客观、有深度。
                """
            elif analysis_type == '折溢价分析':
                prompt = f"""
                请基于以下数据对ETF进行折溢价分析：
                
                基本信息：{json.dumps(etf_data.get('basic_info', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                折溢价数据：{json.dumps(etf_data.get('premium_discount', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                历史数据：{json.dumps(etf_data.get('history', {})[:5], ensure_ascii=False, cls=DateTimeEncoder)}
                
                用户需求：{state['user_query']}
                
                请生成一份专业的ETF折溢价分析报告，包含以下要素：
                1. 当前折溢价水平
                2. 历史折溢价趋势
                3. 折溢价形成原因分析
                4. 套利机会评估
                5. 交易量与折溢价关系
                6. 投资策略建议
                
                使用Markdown格式输出，确保专业、客观、有深度。
                """
            elif analysis_type == '持仓分析':
                prompt = f"""
                请基于以下数据对ETF进行持仓分析：
                
                基本信息：{json.dumps(etf_data.get('basic_info', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                持仓数据：{json.dumps(etf_data.get('portfolio', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                净值数据：{json.dumps(etf_data.get('nav', {})[:5], ensure_ascii=False, cls=DateTimeEncoder)}
                
                用户需求：{state['user_query']}
                
                请生成一份专业的ETF持仓分析报告，包含以下要素：
                1. 行业配置分析
                2. 重仓股分析
                3. 持仓集中度评估
                4. 与标的指数成分股对比
                5. 持仓调整趋势
                6. 投资建议
                
                使用Markdown格式输出，确保专业、客观、有深度。
                """
            else:  # 综合分析
                prompt = f"""
                请基于以下数据对ETF进行综合分析：
                
                基本信息：{json.dumps(etf_data.get('basic_info', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                净值数据：{json.dumps(etf_data.get('nav', {})[:10], ensure_ascii=False, cls=DateTimeEncoder)}
                
                持仓数据：{json.dumps(etf_data.get('portfolio', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                业绩指标：{json.dumps(etf_data.get('performance', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                跟踪误差：{json.dumps(etf_data.get('tracking_error', 0.0), ensure_ascii=False, cls=DateTimeEncoder)}
                
                折溢价数据：{json.dumps(etf_data.get('premium_discount', {}), ensure_ascii=False, cls=DateTimeEncoder)}
                
                历史数据：{json.dumps(etf_data.get('history', {})[:5], ensure_ascii=False, cls=DateTimeEncoder)}
                
                用户需求：{state['user_query']}
                
                请生成一份全面的ETF综合分析报告，包含以下要素：
                1. 核心结论摘要
                2. ETF基本面分析（规模、流动性、费率）
                3. 业绩表现分析（收益情况、波动性）
                4. 跟踪效果分析（跟踪误差、复制策略）
                5. 折溢价分析（当前水平、历史趋势）
                6. 持仓分析（行业配置、重仓股）
                7. 交易策略建议（适合的投资者类型、建议持有期限）
                8. 投资风险提示
                
                使用Markdown格式输出，确保专业、客观、有深度。
                """
            
            # 执行分析
            analysis = self.llm.invoke(prompt)
            
            return {**state, "analysis_result": analysis.content}
                
        except Exception as e:
            self.logger.error(f"分析过程中发生错误: {str(e)}")
            return {**state, "analysis_result": self._format_error_response(str(e))}
    
    def analyze(self, query: str) -> str:
        """执行完整分析流程"""
        try:
            result = self.agent.invoke({"user_query": query})
            return result['analysis_result']
        except Exception as e:
            self.logger.error(f"分析执行失败: {str(e)}")
            return self._format_error_response(str(e))